#include "point_preprocess.h"
#include "pointcloud_tools.h"
#include <fstream>
#include <cassert>

namespace PointPreprocess
{

    void recursive_split(int z_low, int y_low, int x_low,
                         int z_high, int y_high, int x_high,
                         const std::vector<PointCoord> &all_coords,
                         const std::vector<int> &candidate_indices,
                         PartitionResult &result,
                         BlockConfig &bcfg_global,
                         uint64_t &current_pt_addr,
                         uint64_t &current_in_addr,
                         uint64_t &current_rs_addr,
                         const OctreeConfig &config, bool offline_mode)
    {
        if (candidate_indices.empty())
        {
            return;
        }

        if (candidate_indices.size() <= config.block_size)
        {
            // Create a block config
            BlockConfig current_bcfg;
            std::vector<int> current_filter;

            current_bcfg.pt_data.reserve(candidate_indices.size());
            current_bcfg.in_data.reserve(candidate_indices.size() * (config.ich_loop_num + 1));
            current_filter.reserve(candidate_indices.size());

            // Construct a point map
            result.pt_mapping.push_back(candidate_indices);

            // Compute the address in DDR and other parameters
            current_bcfg.BASE_PT_ADDR = current_pt_addr;
            current_bcfg.BASE_IN_ADDR = current_in_addr;
            current_bcfg.BASE_RS_ADDR = current_rs_addr;
            current_bcfg.cfg_cam_valid_num = candidate_indices.size();
            current_bcfg.cfg_total_point_num = candidate_indices.size();

            // Find points that belong to the original range without padding
            int count = 0;
            for (int loop = 0; loop <= config.ich_loop_num; loop++)
            {
                for (int idx : candidate_indices)
                {
                    if (loop == 0)
                    {
                        current_bcfg.pt_data.push_back(bcfg_global.pt_data[idx]);
                        const PointCoord &coord = all_coords[idx];
                        if (isInRange(coord, x_low, x_high, y_low, y_high, z_low, z_high))
                        {
                            current_filter.push_back(count);
                        }
                        count++;
                    }
                    current_bcfg.in_data.push_back(bcfg_global.in_data[idx + loop * bcfg_global.pt_data.size()]);
                }
            }
            if (offline_mode)
            {
                auto align64B = [](uint64_t addr)
                {
                    return (addr + 0xFFF) & ~(uint64_t)(0xFFF);
                };
                current_pt_addr = align64B(current_pt_addr + candidate_indices.size() * sizeof(int64_t));
                current_in_addr = align64B(current_in_addr + candidate_indices.size() * sizeof(int64_t) * (1 + config.ich_loop_num));
                current_rs_addr = align64B(current_rs_addr + candidate_indices.size() * sizeof(int64_t) * (1 + config.och_loop_num));
            }
            // #ifdef SIMULATION
            //             auto align64B = [](uint64_t addr)
            //             {
            //                 return (addr + 0xFFF) & ~(uint64_t)(0xFFF);
            //             };
            //             current_pt_addr = align64B(current_pt_addr + candidate_indices.size() * sizeof(int64_t));
            //             current_in_addr = align64B(current_in_addr + candidate_indices.size() * sizeof(int64_t) * (1 + config.ich_loop_num));
            //             current_rs_addr = align64B(current_rs_addr + candidate_indices.size() * sizeof(int64_t) * (1 + config.och_loop_num));
            // #endif
            result.block_configs.push_back(current_bcfg);
            result.valid_filters.push_back(current_filter);
        }
        else
        {
            // Split the range into 8 sub-blocks
            int z_mid = (z_low + z_high) / 2;
            int y_mid = (y_low + y_high) / 2;
            int x_mid = (x_low + x_high) / 2;

            // Create 8 sub-blocks for recursive searching
            // Padding parameters
            int pad_z = config.search_range[0];
            int pad_y = config.search_range[1];
            int pad_x = config.search_range[2];

            int z_low_pad = z_low - pad_z;
            int y_low_pad = y_low - pad_y;
            int x_low_pad = x_low - pad_x;
            int z_high_pad = z_high + pad_z;
            int y_high_pad = y_high + pad_y;
            int x_high_pad = x_high + pad_x;

            // Initialize 8 sub-region candidates
            std::vector<std::vector<int>> sub_candidates(8);
            for (auto &sub : sub_candidates)
            {
                sub.reserve(candidate_indices.size() / 4 + 1);
            }

            // Note that every candidate region should be padded !
            for (int idx : candidate_indices)
            {
                const PointCoord &coord = all_coords[idx];

                // Every dimension holds three different states for possible overlap after padding
                bool in_z_low = (coord.z >= z_low_pad && coord.z <= z_mid);
                bool in_z_high = (coord.z > z_mid && coord.z <= z_high_pad);
                bool in_z_overlap = (coord.z > z_mid - pad_z && coord.z <= z_mid + pad_z);

                bool in_y_low = (coord.y >= y_low_pad && coord.y <= y_mid);
                bool in_y_high = (coord.y > y_mid && coord.y <= y_high_pad);
                bool in_y_overlap = (coord.y > y_mid - pad_y && coord.y <= y_mid + pad_y);

                bool in_x_low = (coord.x >= x_low_pad && coord.x <= x_mid);
                bool in_x_high = (coord.x > x_mid && coord.x <= x_high_pad);
                bool in_x_overlap = (coord.x > x_mid - pad_x && coord.x <= x_mid + pad_x);

                for (int i = 0; i < 8; i++)
                {
                    bool valid_z = false;
                    bool valid_y = false;
                    bool valid_x = false;

                    if (i / 4 == 0)
                    {
                        valid_z = in_z_low || in_z_overlap;
                    }
                    else
                    {
                        valid_z = in_z_high || in_z_overlap;
                    }

                    if ((i / 2) % 2 == 0)
                    {
                        valid_y = in_y_low || in_y_overlap;
                    }
                    else
                    {
                        valid_y = in_y_high || in_y_overlap;
                    }

                    if (i % 2 == 0)
                    {
                        valid_x = in_x_low || in_x_overlap;
                    }
                    else
                    {
                        valid_x = in_x_high || in_x_overlap;
                    }

                    if (valid_z && valid_y && valid_x)
                    {
                        sub_candidates[i].push_back(idx);
                    }
                }
            }

            for (int i = 0; i < 8; i++)
            {
                if (sub_candidates[i].empty())
                    continue;

                int new_z_low = (i / 4) % 2 == 0 ? z_low : z_mid + 1;
                int new_z_high = (i / 4) % 2 == 0 ? z_mid : z_high;
                int new_y_low = (i / 2) % 2 == 0 ? y_low : y_mid + 1;
                int new_y_high = (i / 2) % 2 == 0 ? y_mid : y_high;
                int new_x_low = i % 2 == 0 ? x_low : x_mid + 1;
                int new_x_high = i % 2 == 0 ? x_mid : x_high;

                recursive_split(new_z_low, new_y_low, new_x_low, new_z_high, new_y_high, new_x_high,
                                all_coords, sub_candidates[i], result, bcfg_global,
                                current_pt_addr, current_in_addr, current_rs_addr, config, offline_mode);
            }
        }
    }

    inline bool partition(BlockConfig &bcfg, const OctreeConfig &config, PartitionResult &result, bool offline_mode)
    {
        if (config.search_range == nullptr || config.spatial_shape == nullptr)
        {
            std::cerr << "Error: search_range or spatial_shape is null" << std::endl;
            return false;
        }

        result.clear();

        // Pre-fetch every point's coord
        std::vector<PointCoord> all_coords;
        all_coords.reserve(bcfg.pt_data.size());

        for (int i = 0; i < bcfg.pt_data.size(); i++)
        {
            all_coords.emplace_back(bcfg.pt_data[i], i);
        }

        // Initialize candidates
        std::vector<int> initial_candidates;
        initial_candidates.reserve(all_coords.size());
        for (int i = 0; i < all_coords.size(); i++)
        {
            initial_candidates.push_back(i);
        }

        // Estimate number of blocks
        int estimated_blocks = std::max(size_t(1), bcfg.pt_data.size() / config.block_size);
        result.pt_mapping.reserve(estimated_blocks);
        result.block_configs.reserve(estimated_blocks);
        result.valid_filters.reserve(estimated_blocks);

        uint16_t vox_x_min = config.spatial_shape[2];
        uint16_t vox_y_min = config.spatial_shape[1];
        uint16_t vox_z_min = config.spatial_shape[0];
        uint16_t vox_x_max = config.spatial_shape[5];
        uint16_t vox_y_max = config.spatial_shape[4];
        uint16_t vox_z_max = config.spatial_shape[3];

        uint64_t current_pt_addr = bcfg.BASE_PT_ADDR;
        uint64_t current_in_addr = bcfg.BASE_IN_ADDR;
        uint64_t current_rs_addr = bcfg.BASE_RS_ADDR;

        recursive_split(vox_z_min, vox_y_min, vox_x_min, vox_z_max, vox_y_max, vox_x_max,
                        all_coords, initial_candidates, result, bcfg,
                        current_pt_addr, current_in_addr, current_rs_addr, config, offline_mode);

        return true;
    }

    inline void mergeFeatures(const PartitionResult &partition_result,
                              int valid_pt_num,
                              int och_loop_num,
                              std::vector<uint64_t> &merged_res,
                              const uint64_t (&bias_cores)[4],
                              uint8_t relu_enable)
    {

        merged_res.assign(valid_pt_num * (1 + och_loop_num), 0);

        for (int i = 0; i < partition_result.block_configs.size(); i++)
        {
            const BlockConfig &bcfg_tmp = partition_result.block_configs[i];
            const std::vector<uint64_t> &rs_tmp = bcfg_tmp.rs_data;
            const std::vector<int> &filter_tmp = partition_result.valid_filters[i];
            const std::vector<int> &pt_mapping_tmp = partition_result.pt_mapping[i];

            for (int k = 0; k <= och_loop_num; k++)
            {
                for (int j = 0; j < filter_tmp.size(); j++)
                {
                    int pt_idx = filter_tmp[j];
                    int pt_mapping_idx = pt_mapping_tmp[pt_idx];

                    uint64_t val = rs_tmp[pt_idx + k * bcfg_tmp.cfg_cam_valid_num];
                    uint64_t bias = bias_cores[k];

                    uint64_t merged_val = 0;
                    // adding bias per channel, 64 channels in total
                    for (int byte_i = 0; byte_i < 8; ++byte_i)
                    {
                        int32_t v = static_cast<int8_t>((val >> (byte_i * 8)) & 0xFF);  // 取出第 byte_i 个字节
                        int32_t b = static_cast<int8_t>((bias >> (byte_i * 8)) & 0xFF); // 取出对应 bias 字节
                        int8_t sum = v + b;                       // 按字节相加（可根据需求决定是否溢出截断）
                        // printf("Byte %d: v = %d, b = %d, sum = %d\n", byte_i, v, b, sum);
                        if(relu_enable == 1 && sum < 0) {
                            sum = 0; // ReLU activation if enabled
                        }
                        // sum = (relu_enable==1) ? std::max(sum, int8_t(0)) : sum; // ReLU activation if enabled
                        merged_val |= (uint64_t(static_cast<uint8_t>(sum)) << (byte_i * 8));
                    }

                    merged_res[pt_mapping_idx + k * valid_pt_num] = merged_val;

                    // merged_res[pt_mapping_idx + k * valid_pt_num] = rs_tmp[pt_idx + k * bcfg_tmp.cfg_cam_valid_num];
                    // adding bias per channel, 64 channels in total
                    // merged_res[pt_mapping_idx + k * valid_pt_num] = rs_tmp[pt_idx + k * bcfg_tmp.cfg_cam_valid_num] + bias_cores[k];
                }
            }
        }

        std::cout << "Merged result done!" << std::endl;
    }

    void exportResults(const std::vector<uint64_t> &merged_res,
                       int valid_pt_num,
                       int och_loop_num,
                       const std::string &filename)
    {

        std::ofstream outfile(filename);
        if (!outfile.is_open())
        {
            std::cerr << "Fail to open " << filename << std::endl;
            return;
        }

        union data64to8
        {
            uint64_t in64;
            int8_t out8[8];
        } converter;

        for (int pt_idx = 0; pt_idx < valid_pt_num; pt_idx++)
        {
            std::string line = "";
            for (int k = 0; k <= och_loop_num; k++)
            {
                uint64_t data = merged_res[pt_idx + k * valid_pt_num];
                converter.in64 = data;

                for (int j = 0; j < 8; j++)
                {
                    if (!line.empty())
                    {
                        line += ",";
                    }
                    line += std::to_string(static_cast<int>(converter.out8[j]));
                }
            }
            outfile << line << std::endl;
        }

        outfile.close();
        std::cout << "MERGED RESULT EXPORT DONE: " << filename << std::endl;
    }

    // For convinence
    void octree_partition(BlockConfig &bcfg,
                          const int block_size,
                          std::vector<std::vector<int>> &pt_mapping,
                          std::vector<BlockConfig> &bcfg_output,
                          std::vector<std::vector<int>> &valid_filter,
                          uint8_t *search_range,
                          uint16_t *spatial_shape,
                          int ich_loop_num,
                          int och_loop_num,
                          bool offline_mode)
    {

        OctreeConfig config;
        config.block_size = block_size;
        config.search_range = search_range;
        config.spatial_shape = spatial_shape;
        config.ich_loop_num = ich_loop_num;
        config.och_loop_num = och_loop_num;

        PartitionResult result;
        if (partition(bcfg, config, result, offline_mode))
        {
            pt_mapping = std::move(result.pt_mapping);
            bcfg_output = std::move(result.block_configs);
            valid_filter = std::move(result.valid_filters);
        }
    }

    void merge_RESULT_feat(const std::vector<BlockConfig> &bcfg_partition,
                           int valid_pt_num,
                           int och_loop_num,
                           const std::vector<std::vector<int>> &pt_mapping,
                           const std::vector<std::vector<int>> &valid_filter,
                           std::vector<uint64_t> &merged_res,
                           const uint64_t (&bias_cores)[4],
                           uint8_t relu_enable)
    {

        PartitionResult temp_result;
        temp_result.block_configs = bcfg_partition;
        temp_result.pt_mapping = pt_mapping;
        temp_result.valid_filters = valid_filter;

        mergeFeatures(temp_result, valid_pt_num, och_loop_num, merged_res, bias_cores, relu_enable);
    }

    void export_merged_results(const std::vector<uint64_t> &merged_res,
                               int valid_pt_num,
                               int och_loop_num,
                               const std::string &filename)
    {
        exportResults(merged_res, valid_pt_num, och_loop_num, filename);
    }

    bool export_blocks_to_files(
        const std::vector<BlockConfig> &block_configs,
        const std::vector<std::vector<int>> &pt_mapping,
        const std::vector<std::vector<int>> &valid_filter,
        const std::string &output_dir,
        const std::vector<uint64_t> &pt_coords,
        const std::vector<uint64_t> &pt_features)
    {
        // Create output directory if it does not exist
        std::string cmd = "mkdir -p " + output_dir;
        system(cmd.c_str());

        for (size_t i = 0; i < block_configs.size(); i++)
        {
            std::string block_dir = output_dir + "/block_" + std::to_string(i);
            cmd = "mkdir -p " + block_dir;
            system(cmd.c_str());

            const std::vector<int> &mapping = pt_mapping[i];
            const std::vector<int> &valid = valid_filter[i];

            std::ofstream coords_file(block_dir + "/coords.txt");
            std::ofstream features_file(block_dir + "/features.txt");
            std::ofstream valid_file(block_dir + "/valid.txt");

            for (size_t j = 0; j < mapping.size(); j++)
            {

                int orig_idx = mapping[j];

                uint64_t coord_data = pt_coords[orig_idx];
                int z = (coord_data >> 24) & 0xFFF;
                int y = (coord_data >> 12) & 0xFFF;
                int x = coord_data & 0xFFF;

                coords_file << orig_idx << " "
                            << x << " "
                            << y << " "
                            << z << " " << std::endl;

                Data64to8 feat_data;
                feat_data.in64 = pt_features[orig_idx];
                features_file << orig_idx << " ";
                for (int k = 0; k < 4; k++)
                {
                    features_file << static_cast<int>(feat_data.out8[k]) << " ";
                }
                features_file << std::endl;
            }

            for (size_t j = 0; j < valid.size(); j++)
            {
                valid_file << valid[j] << std::endl;
            }

            coords_file.close();
            features_file.close();
            valid_file.close();

            std::cout << "Block " << i << " data exported to " << block_dir << std::endl;
        }

        return true;
    }
} // namespace PointPreprocess